import codecs
import os
from prettytable import PrettyTable
import numpy as np

def arrString2floatList(arr_str:str):
    """
    format: array([0.78870674, 0.748     , 0.72161172])
    """
    s = arr_str.split("[")[1].split("]")[0].split(",")
    return [float(val) for val in s]

def ExtractArray(line_str:str):
    """
    BestPerf :  (0.7523035230352304, (array([0.78870674, 0.748     , 0.72161172]), array([0.72287145, 0.83482143, 0.68641115]), array([0.7543554 , 0.78902954, 0.70357143]), array([599, 672, 574])), tensor(1.2319, device='cuda:0'))
    """
    line = line_str.strip('\n').split("(", 1)[1]
    acc_s, line = line.split(',', 1)
    prf1_s, loss_tensor = line.rsplit('),', 1)
    loss_s = loss_tensor.split("tensor(")[1].split(",")[0]
    p, r, f1, cnt = prf1_s[1:].split("),")
    return float(acc_s), float(loss_s), arrString2floatList(p), arrString2floatList(r), arrString2floatList(f1),\
           arrString2floatList(cnt)

def filterFile(filename, maxEpoch=1):
    with open(filename, 'r') as fr:
        with open("./filter.log", 'w') as fw:
            for line in fr:
                if ("####Model Update" in line) and (f"{maxEpoch} | 100" in line):
                    break
                fw.write(line)

def ReadLogFile(filename):
    print("filename : ", filename)
    f1_list = []
    if os.path.exists("./tmp.log"):
        os.system("rm ./tmp.log")
    filterFile(filename, maxEpoch=1)
    os.system(f"cat ./filter.log | grep BestPerf >> tmp.log")
    with codecs.open("./tmp.log", 'r', encoding= u'utf-8',errors='ignore') as fr:
        for line in fr:
                _, _, _, _, f1, _ = ExtractArray(line)
                f1_list.append(f1)
    return f1_list

def exhibit(f1_dic):
    # domain_list = [
    #     'goverment',
    #     'slate',
    #     'telephone',
    #     'travel',
    #     'fiction',
    #     'letters',
    #     'nineeleven',
    #     'oup',
    #     'verbatim',
    #     'facetoface'
    # ]
    for key in f1_dic:
       print(f"{key} : ", np.array(f1_dic[key]).shape)
    table_0 = PrettyTable()
    dic_keys = list(f1_dic.keys())
    dic_keys.sort()
    for key in dic_keys:
        arr = np.array(f1_dic[key])[:, 0]
        table_0.add_column(key, [round(arr.min(), 3), round(arr.mean(), 3), round(arr.max(), 3), round(arr[-1], 3)])
    print("table_0:")
    print(table_0)

    table_1 = PrettyTable()
    for key in dic_keys:
        arr = np.array(f1_dic[key])[:, 1]
        table_1.add_column(key, [round(arr.min(), 3), round(arr.mean(), 3), round(arr.max(), 3), round(arr[-1], 3)])
    print("table_1:")
    print(table_1)

    table_2 = PrettyTable()
    for key in dic_keys:
        arr = np.array(f1_dic[key])[:, 2]
        table_2.add_column(key, [round(arr.min(), 3), round(arr.mean(), 3), round(arr.max(), 3), round(arr[-1], 3)])
    print("table_2:")
    print(table_2)

root_dir = "./"
log_file = [fname for fname in os.listdir(root_dir) if "_FS100.log" in fname]
dic = {}
for fpath in log_file:
    f1list = ReadLogFile(os.path.join(root_dir, fpath))
    dic[fpath.split("_")[1]] =f1list

exhibit(dic)